#include "arch/loongarch.h"
#include "config/type.h"
#include "lib/stdout.h"
#include "lib/string.h"
#include "mm/phylayout.h"
#include "mm/vm.h"

void tlbinit(void) {
	// PS 位为 STLB 页大小的 2 的幂指数，5:0 共 6 位
	// 例：页大小为 16KB：PS=0xE=1110B，2^14=16K
	// 页大小为 4KB：2^12=4096，PS=12=1100B=0xC
	w_csr_stlbps(0xcU);
	// TLB 重填例外专用的页大小值，也设为 4KB
	w_csr_tlbrehi(0xcU);
	// 地址空间表识（ASID，10bits）用于区分不同进程中的统一虚地址
	// 操作系统软件需要为每个进程分配唯一的 ASID
	// w_csr_asid(0x0U); 
}

// TODO: xv6-loongarch 可能未完全实现
void vminit(void) {
	pagetable_t kpgtbl;

	kpgtbl = (pagetable_t)kalloc();
	memset(kpgtbl, 0, PGSIZE);

	// 低半地址空间和高半地址空间实际上就是对虚拟地址空间对半分成了两个部分
	// 虚拟地址最高位为0则表示处于低半地址空间，为1则表示处于高半地址空间
	// 两个部分的地址都用对应部分的 pgd 进行地址转换，互不干扰
	// Linux-loongarch 设定 CSR.PGDH 存放内核进程的 pgd，CSR.PGDL 中存放用户进程的 pgd
	// 硬件上对内核空间和用户空间的页表进行隔离
	pr_debug(vminit, "kpgtbl = %p", kpgtbl);
	proc_mapstacks(kpgtbl);
	w_csr_pgdl((uint64)kpgtbl);

	pr_info(vminit, "tlbinit() ...");
	tlbinit();

	uint64 pwctl0 = (PTEWIDTH << 30) | (DIR2WIDTH << 25) | (DIR2BASE << 20) |
					(DIR1WIDTH << 15) | (DIR1BASE << 10) | (PTWIDTH << 5) |
					(PTBASE << 0);
	uint64 pwctl1 = (DIR4WIDTH << 18) | (DIR4BASE << 12) | (DIR3WIDTH << 6) | (DIR3BASE << 0);
	pr_debug(vminit, "pwcl: %p", pwctl0);
	w_csr_pwcl(pwctl0);
	pr_debug(vminit, "pwch: %p", pwctl1);
	w_csr_pwch(pwctl1);
}

// 64 位的虚拟地址将被分成五个字段：
//	[48...63] -- 必须为零
//	[39...47] -- 9 bits 三级索引
//	[30...38] -- 9 bits 二级索引
//	[21...29] -- 9 bits 一级索引
//	[12...20] -- 9 bits 零级索引
//	[0...11]  -- 12 bits 页内偏移量

// 返回页表 pagetable 中对应于虚拟地址 va 的 PTE 地址。 
// 如果 alloc != 0 就创建需要的页表页。
// 分页方案有四级页表，一个页表页包含 512 个 64 位 PTE。
pte_t *walk(pagetable_t pagetable, uint64 va, int alloc) {
	if (va >= MAXVA)
		panic(walk, "va = %p, MAXVA = %p, va >= MAXVA", va, MAXVA);

	// TODO: 此处 level 循环 3 次过程需要确认
	for (int level = 3; level >= 0; level--) {
		pte_t *pte = &pagetable[PX(level, va)];
		if (*pte & PTE_V) {
			pagetable = (pagetable_t)(PTE2PA(*pte) | DMWIN_MASK);
		} else {
			if (!alloc || (pagetable = (pte_t *)kalloc()) == 0)
				return 0;
			memset(pagetable, 0, PGSIZE);
			*pte = PA2PTE(pagetable) | PTE_V;
		}
	}
	return &pagetable[PX(0, va)];
}

// 查询一个虚拟地址，返回对应的物理地址，如果没有被映射则返回 0
uint64 walkaddr(pagetable_t pagetable, uint64 va) {
	pte_t *pte;
	uint64 pa;

	if (va >= MAXVA)
		return 0;

	pte = walk(pagetable, va, 0);
	if (pte == 0)
		return 0;
	if ((*pte & PTE_V) == 0)
		return 0;
	if ((*pte & PTE_PLV) == 0)
		return 0;
	pa = PTE2PA(*pte);
	return pa;
}

// 为从 va 开始的虚拟地址创建 PTE，这些虚拟地址指向从 pa 开始的物理地址。
// 分配成功时返回 0，如果 walk() 不能分配到所需的页表页，则返回 1。
int mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa,
						 uint64 perm) {
	uint64 a, last;
	pte_t *pte;

	if (size == 0)
		panic(mappages, "size == 0");

	a = PGROUNDDOWN(va);
	last = PGROUNDDOWN(va + size - 1);
	for (;;) {
		if ((pte = walk(pagetable, a, 1)) == 0)
			return -1;
		if (*pte & PTE_V)
			panic(mappages, "*pte & PTE_V == true, remap");
		*pte = PA2PTE(pa) | perm | PTE_V;
		if (a == last)
			break;
		a += PGSIZE;
		pa += PGSIZE;
	}
	return 0;
}

// 移除从 va 开始的 npages 映射，va 必须是页对齐的。
// 这些映射必须存在。可以选择释放物理内存。
void uvmunmap(pagetable_t pagetable, uint64 va, uint64 npages, int do_free) {
	uint64 a;
	pte_t *pte;

	if ((va % PGSIZE) != 0)
		panic(uvmunmap, "(va %% PGSIZE) != 0, not aligned");

	for (a = va; a < va + npages * PGSIZE; a += PGSIZE) {
		if ((pte = walk(pagetable, a, 0)) == 0)
			panic(uvmunmap, "(pte = walk(pagetable, a, 0)) == 0");
		if ((*pte & PTE_V) == 0)
			panic(uvmunmap, "(*pte & PTE_V) == 0");
		if (PTE_FLAGS(*pte) == PTE_V)
			panic(uvmunmap, "PTE_FLAGS(*pte) == PTE_V, not a leaf");
		if (do_free) {
			uint64 pa = PTE2PA(*pte);
			kfree((void *)(pa | DMWIN_MASK));
		}
		*pte = 0;
	}
}

// 创建一个空的用户页表，如果内存不足，则返回 0。
pagetable_t uvmcreate() {
	pagetable_t pagetable;
	pagetable = (pagetable_t)kalloc();
	if (pagetable == 0)
		return 0;
	memset(pagetable, 0, PGSIZE);
	return pagetable;
}

// 将用户的 initcode 加载到 pagetable 的地址 0，用于第一个进程。
// sz 必须小于一个页面。
void uvminit(pagetable_t pagetable, uchar *src, uint sz) {
	char *mem;

	if (sz >= PGSIZE)
		panic(uvminit, "sz >= PGSIZE, more than a page");
	mem = kalloc();
	memset(mem, 0, PGSIZE);
	mappages(pagetable, 0, PGSIZE, (uint64)mem, PTE_P | PTE_W | PTE_PLV | PTE_MAT);
	memmove(mem, src, sz);
} // TODO: xv6-loongarch 未完全实现

// 删除用户页，使进程大小从oldsz变为newsz。
// oldsz和newsz不需要页对齐，newsz也不需要小于oldsz，oldsz可以大于实际进程大小。 
// 返回新的进程大小。
uint64 uvmdealloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz) {
	if (newsz >= oldsz)
		return oldsz;

	if (PGROUNDUP(newsz) < PGROUNDUP(oldsz)) {
		int npages = (PGROUNDUP(oldsz) - PGROUNDUP(newsz)) / PGSIZE;
		uvmunmap(pagetable, PGROUNDUP(newsz), npages, 1);
	}

	return newsz;
}

// 分配 PTE 和物理内存，使进程从 oldsz 增长到 newsz，不需要页面对齐。
// 错误时返回新的大小或 0。
uint64 uvmalloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz) {
	char *mem;
	uint64 a;

	if (newsz < oldsz)
		return oldsz;

	oldsz = PGROUNDUP(oldsz);
	for (a = oldsz; a < newsz; a += PGSIZE) {
		mem = kalloc();
		if (mem == 0) {
			uvmdealloc(pagetable, a, oldsz);
			return 0;
		}
		memset(mem, 0, PGSIZE);
		if (mappages(pagetable, a, PGSIZE, (uint64)mem,
								 PTE_P | PTE_W | PTE_PLV | PTE_MAT | PTE_D) != 0) {
			kfree(mem);
			uvmdealloc(pagetable, a, oldsz);
			return 0;
		}
	}
	return newsz;
}

// 递归释放页表页面。
// 所有的叶子映射必须已经被移除。
void freewalk(pagetable_t pagetable) {
	// there are 2^9 = 512 PTEs in a page table.
	for (int i = 0; i < 512; i++) {
		pte_t pte = pagetable[i];
		if ((pte & PTE_V) && (PTE_FLAGS(pte) == PTE_V)) {
			// this PTE points to a lower-level page table.
			uint64 child = (PTE2PA(pte) | DMWIN_MASK);
			freewalk((pagetable_t)child);
			pagetable[i] = 0;
		} else if (pte & PTE_V) {
			panic(freewalk, "pte & PTE_V, leaf");
		}
	}
	kfree((void *)pagetable);
}

// 释放用户内存页，然后释放页表页。
void uvmfree(pagetable_t pagetable, uint64 sz) {
	if (sz > 0)
		uvmunmap(pagetable, 0, PGROUNDUP(sz) / PGSIZE, 1);
	freewalk(pagetable);
}

// 给出一个父进程的页表，将其内存复制到子进程的页表中（同时复制页表和物理内存）。
// 成功时返回0，失败时返回1。
// 失败时释放所有分配的页。
int uvmcopy(pagetable_t old, pagetable_t new, uint64 sz) {
	pte_t *pte;
	uint64 pa, i;
	uint flags;
	char *mem;

	for (i = 0; i < sz; i += PGSIZE) {
		if ((pte = walk(old, i, 0)) == 0)
			panic(uvmcopy, "(pte = walk(old, i, 0)) == 0, pte should exist");
		if ((*pte & PTE_V) == 0)
			panic(uvmcopy, "(*pte & PTE_V) == 0, page not present");
		pa = PTE2PA(*pte);
		flags = PTE_FLAGS(*pte);
		if ((mem = kalloc()) == 0)
			goto err;
		memmove(mem, (char *)(pa | DMWIN_MASK), PGSIZE);
		if (mappages(new, i, PGSIZE, (uint64)mem, flags) != 0) {
			kfree(mem);
			goto err;
		}
	}
	return 0;

err:
	uvmunmap(new, 0, i / PGSIZE, 1);
	return -1;
}

// 标记一个PTE为无效的用户访问，该函数被exec用于用户堆栈保卫页。
void uvmclear(pagetable_t pagetable, uint64 va) {
	pte_t *pte;

	pte = walk(pagetable, va, 0);
	if (pte == 0)
		panic(uvmclear, "pte == 0");
	*pte &= ~PTE_PLV;
}

// 从内核拷贝到用户，从src复制len字节到指定页表中的虚拟地址dstva。
// 成功时返回0，错误时返回1。
int copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len) {
	uint64 n, va0, pa0;

	while (len > 0) {
		va0 = PGROUNDDOWN(dstva);
		pa0 = walkaddr(pagetable, va0);
		if (pa0 == 0)
			return -1;
		n = PGSIZE - (dstva - va0);
		if (n > len)
			n = len;
		memmove((void *)((pa0 + (dstva - va0)) | DMWIN_MASK), src, n);

		len -= n;
		src += n;
		dstva = va0 + PGSIZE;
	}
	return 0;
} // TODO: xv6-loongarch 可能未完全实现

// 从用户拷贝到内核。从给定页表中的虚拟地址srcva复制len字节到dst。
// 成功时返回0，错误时返回1。
int copyin(pagetable_t pagetable, char *dst, uint64 srcva, uint64 len) {
	uint64 n, va0, pa0;

	while (len > 0) {
		va0 = PGROUNDDOWN(srcva);
		pa0 = walkaddr(pagetable, va0);
		if (pa0 == 0)
			return -1;
		n = PGSIZE - (srcva - va0);
		if (n > len)
			n = len;
		memmove(dst, (void *)((pa0 + (srcva - va0)) | DMWIN_MASK), n);

		len -= n;
		dst += n;
		srcva = va0 + PGSIZE;
	}
	return 0;
} // TODO: xv6-loongarch 可能未完全实现

// 从用户到内核拷贝一个空尾的字符串。
// 从给定页表的虚拟地址 srcva 拷贝字节到 dst、
// 直到出现'\0'，或长度达到max。
// 成功时返回0，错误时返回1。
int copyinstr(pagetable_t pagetable, char *dst, uint64 srcva, uint64 max) {
	uint64 n, va0, pa0;
	int got_null = 0;

	while (got_null == 0 && max > 0) {
		va0 = PGROUNDDOWN(srcva);
		pa0 = walkaddr(pagetable, va0);
		if (pa0 == 0)
			return -1;
		n = PGSIZE - (srcva - va0);
		if (n > max)
			n = max;

		char *p = (char *)((pa0 + (srcva - va0)) | DMWIN_MASK);
		while (n > 0) {
			if (*p == '\0') {
				*dst = '\0';
				got_null = 1;
				break;
			} else {
				*dst = *p;
			}
			--n;
			--max;
			p++;
			dst++;
		}

		srcva = va0 + PGSIZE;
	}
	if (got_null) {
		return 0;
	} else {
		return -1;
	}
} // TODO: xv6-loongarch 可能未完全实现